import sys
sys.path.append("../../")
from datareader import *
import json

NLI_domain_list = list(NLI_domain_map.keys())
few_shot_max = 100
with open("../../../multinli_1.0/multinli_1.0_dev_mismatched.jsonl") as fr:
    lines = [line for line in fr if json.loads(line.strip("\n"))['gold_label'] != '-']
    genres = list(set([json.loads(line.strip("\n"))['genre'] for line in lines]))
    grouped_lines = [[line for line in lines if json.loads(line.strip("\n"))['genre']==g] for g in genres]
    for idx, g in enumerate(genres):
        g_lines = random.sample(grouped_lines[idx], len(grouped_lines[idx]))
        with open(f"../../../multinli_1.0/fewshot_Domain_{g}.jsonl", "w") as fw:
            fw.write("".join(g_lines[:few_shot_max]))
        with open(f"../../../multinli_1.0/Domain_{g}.jsonl", "w") as fw:
            fw.write("".join(g_lines[few_shot_max:]))

with open("../../../multinli_1.0/multinli_1.0_dev_matched.jsonl") as fr:
    lines = [line for line in fr if json.loads(line.strip("\n"))['gold_label'] != '-']
    genres = list(set([json.loads(line.strip("\n"))['genre'] for line in lines]))
    grouped_lines = [[line for line in lines if json.loads(line.strip("\n"))['genre']==g] for g in genres]
    for idx, g in enumerate(genres):
        g_lines = random.sample(grouped_lines[idx], len(grouped_lines[idx]))
        with open(f"../../../multinli_1.0/fewshot_Domain_{g}.jsonl", "w") as fw:
            fw.write("".join(g_lines[:few_shot_max]))
        with open(f"../../../multinli_1.0/Domain_{g}.jsonl", "w") as fw:
            fw.write("".join(g_lines[few_shot_max:]))